package com.example.graph;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * 给你一个 n 个点的带权无向连通图，节点编号为 0 到 n-1 ，同时还有一个数组 edges ，其中 edges[i] = [fromi, toi, we
 * ighti] 表示在 fromi 和 toi 节点之间有一条带权无向边。
 * 最小生成树 (MST) 是给定图中边的一个子集，它连接了所有节点且没有环，而且这些边的权值和最小。
 * 请你找到给定图中最小生成树的所有关键边和伪关键边。
 * 如果从图中删去某条边，会导致最小生成树的权值和增加，那么我们就说它是一条关键边。
 * 伪关键边则是可能会出现在某些最小生成树中但不会出现在所有最小生成树中的边。
 * <p>
 * 请注意，你可以分别以任意顺序返回关键边的下标和伪关键边的下标。
 * <p>
 * 示例 1：
 * 输入：n = 5, edges = [[0,1,1],[1,2,1],[2,3,2],[0,3,2],[0,4,3],[3,4,3],[1,4,6]]
 * 输出：[[0,1],[2,3,4,5]](索引即是边数组的索引)
 * <p>
 * 注意到第 0 条边和第 1 条边出现在了所有最小生成树中，所以它们是关键边，我们将这两个下标作为输出的第一个列表。
 * 边 2，3，4 和 5 是所有 MST 的剩余边，所以它们是伪关键边。我们将它们作为输出的第二个列表。
 * <p>
 * 示例 2 ：
 * <p>
 * 输入：n = 4, edges = [[0,1,1],[1,2,1],[2,3,1],[0,3,1]]
 * 输出：[[],[0,1,2,3]]
 * <p>
 * 提示：
 * 2 <= n <= 100
 * 1 <= edges.length <= min(200, n * (n - 1) / 2)
 * edges[i].length == 3
 * 0 <= fromi < toi < n
 * 1 <= weighti <= 1000
 * 所有 (fromi, toi) 数对都是互不相同的。
 */
public class Leetcode1489_FindCriticalAndPseudoCriticalEdges {

    public static void main(String[] args) {

        int[][] edges;
        int n = 5;
        List<int[]> mstEdges;

//        edges = new int[][]{
//                {0, 1, 1}, {1, 2, 1}, {2, 3, 2}, {0, 3, 2},
//                {0, 4, 3}, {3, 4, 3}, {1, 4, 6}};
        edges = new int[][]{
                {0, 1, 1}, {1, 2, 1}, {2, 3, 1}, {0, 3, 1}
        };
//        edges = new int[][]{
//                {0, 1, 6}, {0, 2, 1}, {0, 3, 5},
//                {1, 2, 5}, {1, 4, 3},
//                {2, 3, 5}, {2, 4, 6}, {2, 5, 4},
//                {3, 5, 2}, {4, 5, 6}};

        List<List<Integer>> res = new Solution().findCriticalAndPseudoCriticalEdges(n, edges);

        res.forEach(e -> {
            e.forEach(x -> {
                System.out.print(x + ",");
            });
            System.out.println();
        });


    }

    static class Solution {

        public List<List<Integer>> findCriticalAndPseudoCriticalEdges(int n, int[][] edges) {
            return findCriticalAndPseudoCriticalEdges2(n, edges);
        }

        /**
         * 首先根据经典的kruskal算法得到最小生成树的权值和
         * 然后求关键边和伪关键边
         * 遍历所有的边执行如下操作：
         * 在删掉某一条后边再重新运行kruskal算法求最小生成树的权值和
         * 如果权值和变大则是关键边
         * <p>
         * 在加入某一条边后再重新运行kruskal算法求最小生成树的权值和
         * 如果不变则说明是伪关键边
         *
         * @param n
         * @param edges
         * @return
         */
        private int[][] edgesCopy; // kruskal要使用

        public List<List<Integer>> findCriticalAndPseudoCriticalEdges1(int n, int[][] edges) {
            List<List<Integer>> res = new ArrayList<>();
            List<Integer> keyEdges = new ArrayList<>(); // 关键边集合
            List<Integer> pseuEdges = new ArrayList<>(); // 伪关键边集合
            res.add(keyEdges);
            res.add(pseuEdges);
            edgesCopy = edges.clone();

            Arrays.sort(edges, (e1, e2) -> e1[2] - e2[2]);
            int mstWeightSum = kruskal(edges, n, -1, false);
//            for (int i = 0; i < edges.length; i++) {
//                for (int j = 0; j < edges[0].length; j++) {
//                    System.out.print(edges[i][j] + ", ");
//                }
//                System.out.println();
//            }

            for (int i = 0; i < edgesCopy.length; i++) {
                if (kruskal(edges, n, i, true) > mstWeightSum) {
                    keyEdges.add(i);
                } else if (kruskal(edges, n, i, false) == mstWeightSum) {
                    pseuEdges.add(i);
                }
            }

            return res;
        }

        /**
         * prim算法求MST
         * 整体逻辑通法一
         * 不同的就是求生成树使用Prim算法
         *
         * @param n
         * @param edges
         * @return
         */
        public List<List<Integer>> findCriticalAndPseudoCriticalEdges2(int n, int[][] edges) {
            List<List<Integer>> res = new ArrayList<>();
            List<Integer> keyEdges = new ArrayList<>(); // 关键边集合
            List<Integer> pseuEdges = new ArrayList<>(); // 伪关键边集合
            res.add(keyEdges);
            res.add(pseuEdges);
            makeAdjMatrix(edges, n);

            int mstWeightSum = primByAdjMatrix(edges, n, -1, false);

            for (int i = 0; i < edges.length; i++) {
                if (primByAdjMatrix(edges, n, i, true) > mstWeightSum) {
                    keyEdges.add(i);
                } else if (primByAdjMatrix(edges, n, i, false) == mstWeightSum) {
                    pseuEdges.add(i);
                }
            }

            return res;
        }

        /**
         * kruskal算法求最小生成树
         *
         * @param sortedEdges
         * @param n
         * @param index       要预先加入最小生成树的边或者去掉边在初始数组的索引(-1表示既不加入边也不删掉边)
         * @param giveOrTake  提前给定一条边时为false，提前去掉一条边为true
         * @return
         */
        private int kruskal(int[][] sortedEdges, int n, int index, boolean giveOrTake) {
            int mstWeightSum = 0, mstEdgeCount = 0;
            // 并查集初始化
            init(n);

            // 伪关键边处理，先将边加入到MST中(即 将边的两个顶点合并到一个集合中)
            if (!giveOrTake && index >= 0) {
                union(edgesCopy[index][1], edgesCopy[index][0]); // 完成预先加入边的两个顶点的合并
                mstWeightSum += edgesCopy[index][2];
                mstEdgeCount++;
            }

            for (int i = 0; i < sortedEdges.length && mstEdgeCount < n - 1; i++) {
                if (giveOrTake && sortedEdges[i] == edgesCopy[index]) {
                    continue; // 跳过就想到于删除了该边
                }

                if (union(sortedEdges[i][0], sortedEdges[i][1])) { // 如果需要合并说明两点不连通
                    mstWeightSum += sortedEdges[i][2];
                    mstEdgeCount++;
                }
            }

            if (mstEdgeCount == n - 1) // 正常合并了n-1次
                return mstWeightSum;

            return Integer.MAX_VALUE;// 去掉关键边之后，可能发生不连通的情况就返回一个非法值
        }

        private int primByAdjMatrix(int[][] edges, int n, int index, boolean giveOrTake) {
            // mst权值和，边数量
            int mstWeightSum = 0, mstVertexCount = 0;
            // 集合A，临时存放顶点
            boolean[] visited = new boolean[n];
            // 伪关键边处理
            if (!giveOrTake && index >= 0) {
                visited[edges[index][0]] = true;
                visited[edges[index][1]] = true;
                mstWeightSum += edges[index][2];
                mstVertexCount = 2;
            } else { // 默认
                visited[0] = true; // 0节点加入A
                mstVertexCount++;
            }

            while (mstVertexCount < n) {
                /* 选择这些顶点权值最小的边加入MST */
                int x = -1, y = -1; // 最小权值边的两个顶点
                int minWeight = Integer.MAX_VALUE;
                // 遍历集合 A 中所有顶点，找到最小权值的边
                for (int i = 0; i < n; i++) { // 所有顶点
                    if (visited[i]) { // 只需要已被标记的元素
                        int[] neighbors = matrix[i];
                        for (int j = 0; j < n; j++) {
                             // 0表示没有边
                            if (neighbors[j] == 0) continue;
                             // 从 i 到 j 的边已经被加入mst
                            if (visited[j]) continue;
                            // 关键边处理
                            if (giveOrTake) {
                                int a = edges[index][0]; // 关键边的顶点
                                int b = edges[index][1]; // 关键边的顶点
                                // 两端顶点相同则为关键边
                                if ((a == i && b == j) || (a == j && b == i))
                                    continue;
                            }
                            /* 记录较小权值 */
                            if (neighbors[j] < minWeight) {
                                minWeight = neighbors[j]; // 记录最小权重
                                x = i; // 记录边的两个顶点
                                y = j;
                            }
                        }
                    }
                }
                // 可能会发生去掉关键边之后，不连通的情况
                if (x != -1 && y != -1) {
//                    visited[x] = true;
                    visited[y] = true;
                    mstVertexCount++;
                    mstWeightSum += minWeight; // 累计权重
                } else { // 跳过的边，是连接了两个岛屿的关键边
                    return Integer.MAX_VALUE;
                }
            }
            return mstWeightSum;
        }

        private int[][] matrix;
        private void makeAdjMatrix(int[][] edges, int n) {
            // 创建 n 个节点的邻接矩阵
            matrix = new int[n][n];
            for (int i = 0; i < edges.length; ++i) {
                matrix[edges[i][0]][edges[i][1]] = edges[i][2];
                matrix[edges[i][1]][edges[i][0]] = edges[i][2];
            }
        }

        /*******************并查集操作******************/
        private int[] parent;
        private int[] rank;

        // 初始化并查集
        private void init(int n) {
            parent = new int[n]; // 存每个节点的父节点
            rank = new int[n];// 存每个节点的秩

            for (int i = 0; i < n; i++) {
                parent[i] = i;
                rank[i] = 1;
            }

        }

        // 查找指定节点的父节点,同时完成路径压缩
        private int find(int node) {
            while (parent[node] != node) {
                node = parent[node];
                parent[node] = parent[parent[node]];
            }

            return parent[node];
        }

        // 合并两个集合,不需要合并返回false,需要合并的话,按秩合并后返回true
        private boolean union(int set1, int set2) {
            int root1 = find(set1);
            int root2 = find(set2);
            if (root1 != root2) {
                if (rank[root1] > rank[root2]) {
                    parent[root2] = root1;

                } else if (rank[root1] < rank[root2]) {
                    parent[root1] = root2;
                } else {
                    parent[root2] = root1;
                    rank[root1] += 1;
                }
                return true;
            }
            return false;
        }
        /*********************并查集操作end**************************/
    }
}
